﻿using System.Collections.Generic;
using System.Linq;
using Algoritmia.Utils;

namespace Algoritmia.Problems.Puzzles.NQueens
{
    public class NQueensStateSpace5 : NQueensStateSpace4
    {
        public new class State : NQueensStateSpace4.State
        {
            public State(int N, State prev = null, int row = -1)
                : base(N, prev, row)
            {
            }

            public override IEnumerable<int> FreeRows()
            {
                var col = N - 1;
                if ( col == 0)
                {
                    for(int row = 0; row < (N+1)/2; row++)
                    {
                        yield return row;
                    }
                }
                else
                {
                    foreach(var row in AvailableRows)
                    {
                        if ( CellIsFree(col,row))
                        {
                            yield return row;
                        }
                    }
                }
            }
        }

        public NQueensStateSpace5(int N)
            : base(N)
        {
        }

    }

    public class SpecialNQueensEnumerator : NQueensEnumerator4
    {
        public override IEnumerable<DataStructures.Lists.IList<int>> Enumerate(NQueensStateSpace4 space)
        {
            if ( space is NQueensStateSpace5)
            {
                foreach (var solution in base.Enumerate(space))
                {
                    yield return solution;
                    if (solution[0] < space.N / 2)
                    {
                        yield return solution.Select(i => space.N - i - 1).ToAlgoritmiaList();
                    }
                }
            }
            else
            {
                foreach(var solution in base.Enumerate(space))
                {
                    yield return solution;
                }
            }
            
        }
    }
}
